package vpc.ext.stacksize;

import vpc.tir.*;
import vpc.tir.opt.TIRCallShape;
import vpc.tir.opt.DepthFirstOpVisitor;
import vpc.tir.expr.Operator;
import vpc.tir.TIRConst;
import vpc.core.Program;
import vpc.core.types.TypeRef;
import vpc.core.types.Type;
import vpc.core.decl.Method.Temporary;
import vpc.core.decl.Method;
import vpc.core.decl.Method.Temporary;
import vpc.core.ProgramDecl.EntryPoint;
import vpc.core.types.Type;
import vpc.hil.*;
import vpc.hil.DeviceDecl.GetRegister;
import vpc.hil.DeviceDecl.SetRegister;
import vpc.ext.stacksize.util.Debug;
import vpc.ext.stacksize.util.DataHandler;
import vpc.ext.stacksize.util.VstackException;

import java.util.*;

public class GraphVisitor extends
		DepthFirstOpVisitor<TIRExpr, GraphVisitor.Context> {

	private static Debug.Flag debugFlag = Debug.Flag.GRAPHVISITOR;

	public class Context {
		public Node root;

		public Node parent;

		public Context(Node root) {
			this.root = root;
			this.parent = root;
		}
	}

	public final Context context;

	public final List<Node> subRoots;

	private final Program p;

	public GraphVisitor(Program p) {
		super(null);
		Method m = p.programDecl.mainEntry.method;
		Node root = Node.constructMethodNode(m);
		context = new Context(root);
		subRoots = new ArrayList<Node>();
		subRoots.add(root);
		this.p = p;
	}

	public void setParentByName(String name) {
		this.context.parent = null;

		for (Node n : this.subRoots) {
			if (n.name.equals(name)) {
				this.context.parent = n;
				break;
			}

			Node result = n.find(name, new HashSet<Node>());

			if (result != null) {
				this.context.parent = result;
				break;
			}
		}
	}

	/**
	 * Visit method for TIRExpr that is returned as the result of the getBody()
	 * method.
	 * 
	 * @param body
	 *            TIRExpr representing the body of a method.
	 * @return
	 */
	public TIRExpr visitBody(TIRExpr body) {
		return body.accept(this, this.context);
	}

	private void printInterrupts(Register r, Set<Interrupt> list) {
		Debug.println("REGISTER : " + r.getName(), debugFlag);
		Debug.print("( ", debugFlag);
		if (list != null) {
			for (Interrupt i : list) {
				Debug.println(i.getName() + ",", debugFlag);
			}
		}
		Debug.println(")", debugFlag);

		return;
	}

	/**
	 * In this method, we visit a known SetRegister operation in the TIR Expression tree.
	 * A SetRegister is the only operation in Virgil that can change the state
	 * of the interrupt mask. This method either statically determines the set of enabled/
	 * disabled interrupts (if the argument is a raw, constant value) or it conservatively
	 * assumes that all interrupts are enabled by the SetRegister operation.
	 * @param setreg The TIRExpr object that represents writing to a register
	 * @param operands The sub-tree of TIRExpr(s) that represent the values to be written
	 * @return Null
	 */
	public TIRExpr visit(SetRegister setreg, TIRExpr[] operands) {

		//Get the register that is written to
		Register register = setreg.register;

		//Check to see if the target register belongs to the target device.
		if (InterruptMask.getInterrupts(register) == null) {
			//In this case, we are setting a register that is NOT
			//specified in the program's target device HIL file.
			//Most likey, this is due to an operation occuring in a library.
			//In any case, there ARE cases of SetRegister expressions that
			//do not affect a register belonging to the target device.
			//Hence, we cannot determine the set of affected interrupts.
			//TODO: is the information for these registers located anywhere else
			//in the compiler? Or is it determined at link-time?
			return null;
		}
		//Else, the register is part of the target device.

		InterruptMask newMask = null;
		boolean maskConstructed = false;
		while (maskConstructed == false) {
			try {
				newMask = new InterruptMask();
				maskConstructed = true;
			} catch (VstackException.ClassInitializationException e) {
				InterruptMask.init(this.p.targetDevice);
			}
		}
		assert (newMask != null);

		// a SetReg operation is a UNARY operator; hence, it should only accept
		// one sub-expression.
		assert (operands.length == 1);
		// Get the sole operand
		TIRExpr operand = operands[0];

		// If the operation takes a raw, constant argument that we can
		// analyze at compile-time:
		if (operand instanceof TIRConst.Value) {
			TIRConst.Value v = (TIRConst.Value) operand;

			//Given the constant value (TIRConst.Value), get the list
			//of interrupts that can can be enabled by setting the register
			Set<Interrupt> enabledInterrupts = InterruptMask.getEnabled(
					register, v.getValue());

			for (Interrupt i : enabledInterrupts) {
				newMask.enableInterrupt(i);
			}
						
		} else {
			//Else, the SetRegister operation takes another TIRExpr tree as an
			//argument. We cannot statically determine what the effects of setting
			//the register will do to the interrupt mask. Conservatively assume
			//that the SetRegister operation will enable ALL interrupts at this point.
			newMask.enabledAll();
		}
		
		//At this point, we have created the new mask. It either reflects the exact
		//enabled interrupts, or conservativley enableds ALL interrupts.
		
		//Create new nodes here.

		return null;
	}

	public TIRExpr visitOperator(Operator o, TIRExpr[] operands) {

		if (o instanceof SetRegister) {
			this.visit((SetRegister) o, operands);
		}
		return null;
	}

	/**
	 * Visit method for TIROperator type. Visit each sub-operand.
	 */
	public TIRExpr visit(TIROperator o, Context c) {

		for (TIRExpr e : o.operands)
			e.accept(this, c);

		this.visitOperator(o.operator, o.operands);

		return null;
	}

	/**
	 * Visit method for TIRConst.Symbol type. Represents a String, essentially.
	 * {@link vpc.tir.TIRConst.Symbol}
	 */
	public TIRExpr visit(TIRConst.Symbol s, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRConst.Value type. {@link vpc.tir.TIRConst.Value}
	 */
	public TIRExpr visit(TIRConst.Value v, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRLocal.Get type. {@link vpc.tir.TIRLocal.Get}
	 */
	public TIRExpr visit(TIRLocal.Get g, Context c) {

		return null;
	}

	/**
	 * Visit method for TIRLocal.Set type. {@link vpc.tir.TIRLocal.Set}
	 */
	public TIRExpr visit(TIRLocal.Set l, Context c) {
		l.value.accept(this, c);
		return null;
	}

	/**
	 * Visit method for TIRCall type. NOTE: extremely important for call graph
	 * construction.
	 * 
	 * @author Philip Russell
	 * @author Micah Wendell {@link vpc.tir.TIRCall}
	 */
	public TIRExpr visit(TIRCall b, Context c) {

		Debug.Flag localFlag = Debug.Flag.VISIT_TIRCALL;

		// Need to visit the argument TIRExpr(s) as well
		for (TIRExpr arg : b.arguments) {
			arg.accept(this, this.context);
		}

		TIRCallShape.CallShape cs = TIRCallShape.match(p, b);

		List<Method> methodsInvoked = new ArrayList<Method>();

		if (cs instanceof TIRCallShape.Direct) {
			methodsInvoked.add(((TIRCallShape.Direct) cs).method);
		} else if (cs instanceof TIRCallShape.Dynamic) {

			for (Method m : p.closure.methods) {
				TypeRef[] params = m.getArgumentTypes();
				int numArgs = b.arguments.length;

				if (params.length != numArgs)
					continue;

				if (b.getType() != m.getType())
					continue;

				for (int i = 0; i < numArgs; i++) {
					if (params[i].getType() != b.arguments[i].getType())
						continue;
				}

				methodsInvoked.add(m);
			}
		} else { // Virtual

			// FUTURE TODO: add Methods to the methodsInvoked list using CHA

			Method toCall = ((TIRCallShape.Virtual) cs).method;
			TypeRef[] toCallParams = toCall.getArgumentTypes();
			int numArgs = toCallParams.length;

			for (Method m : p.closure.methods) {
				TypeRef[] params = m.getArgumentTypes();

				if (!m.getName().equals(toCall.getName()))
					continue;

				if (params.length != numArgs)
					continue;

				if (b.getType() != m.getType())
					continue;

				for (int i = 0; i < numArgs; i++) {
					if (params[i].getType() != toCallParams[i].getType())
						continue;
				}

				methodsInvoked.add(m);
			}
		}
		for (Method m : methodsInvoked) {

			Node child = null;
			for (Node n : this.subRoots) {

				Debug.println("\tn.id = " + n.name + "; m.getFullName() = "
						+ m.getFullName(), localFlag);

				if (n.name.equals(m.getFullName())) {

					Debug.println("\tFound matching node in subroots, id = "
							+ n.name, localFlag);

					child = n;
					if (n != context.root && n != c.parent.root) {
						this.subRoots.remove(n);
						n.changeRoot(c.parent.root, new HashSet<String>());
					}
					break;
				}

				child = n.find(m.getFullName(), new HashSet<Node>());

				Debug.println("\tResult of find(): " + child, localFlag);

				if (child != null)
					break;
			}

			if (child == null) {
				child = new Node(m.getFullName(), DataHandler.getStackPush(m),
						c.parent.root);

				Debug.println("\tCreating " + m.getFullName() + " in subgraph",
						localFlag);
			}

			c.parent.addEdge(child);

			Debug.println("\tAdded edge to " + m.getFullName(), localFlag);
		}

		return null;
	}

	/**
	 * Generic visit method for TIRExpr(s) not listed above.
	 * {@link vpc.tir.TIRExpr}
	 */
	public TIRExpr visit(TIRExpr e, Context c) {

		Debug.println("WARNING: visited unknown TIRExpr, CLASS = "
				+ e.getClass().toString(), debugFlag);
		return null;
	}
}
